{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b74e860a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b7a1efb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SelfAttention(nn.Module):\n",
    "    def __init__(self, embed_size :int, heads :int):\n",
    "        super(SelfAttention, self).__init__()\n",
    "        self.embed_size = embed_size\n",
    "        self.heads = heads\n",
    "        # 向下取整\n",
    "        self.head_dim = embed_size // heads\n",
    "\n",
    "        # 判定Dimension正确\n",
    "        assert (\n",
    "            self.head_dim * heads == embed_size\n",
    "        ), \"Embedding size needs to be divisible by heads\"\n",
    "\n",
    "        self.values = nn.Linear(embed_size, embed_size)\n",
    "        self.keys = nn.Linear(embed_size, embed_size)\n",
    "        self.queries = nn.Linear(embed_size, embed_size)\n",
    "        self.fc_out = nn.Linear(embed_size, embed_size)\n",
    "\n",
    "    def forward(self, values, keys, query, mask):\n",
    "        # Get number of training examples\n",
    "        N = query.shape[0]\n",
    "\n",
    "        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]\n",
    "\n",
    "        values = self.values(values)  # (N, value_len, embed_size)\n",
    "        keys = self.keys(keys)  # (N, key_len, embed_size)\n",
    "        queries = self.queries(query)  # (N, query_len, embed_size)\n",
    "\n",
    "        # Split the embedding into self.heads different pieces\n",
    "        values = values.reshape(N, value_len, self.heads, self.head_dim)\n",
    "        keys = keys.reshape(N, key_len, self.heads, self.head_dim)\n",
    "        queries = queries.reshape(N, query_len, self.heads, self.head_dim)\n",
    "\n",
    "        # Einsum does matrix mult. for query*keys for each training example\n",
    "        # with every other training example, don't be confused by einsum\n",
    "        # it's just how I like doing matrix multiplication & bmm\n",
    "\n",
    "        energy = torch.einsum(\"nqhd,nkhd->nhqk\", [queries, keys])\n",
    "        # queries shape: (N, query_len, heads, heads_dim),\n",
    "        # keys shape: (N, key_len, heads, heads_dim)\n",
    "        # energy: (N, heads, query_len, key_len)\n",
    "\n",
    "        # Mask padded indices so their weights become 0\n",
    "        if mask is not None:\n",
    "            energy = energy.masked_fill(mask == 0, float(\"-1e20\"))\n",
    "\n",
    "        # Normalize energy values similarly to seq2seq + attention\n",
    "        # so that they sum to 1. Also divide by scaling factor for\n",
    "        # better stability\n",
    "        attention = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=3)\n",
    "        # attention shape: (N, heads, query_len, key_len)\n",
    "\n",
    "        out = torch.einsum(\"nhql,nlhd->nqhd\", [attention, values]).reshape(\n",
    "            N, query_len, self.heads * self.head_dim\n",
    "        )\n",
    "        # attention shape: (N, heads, query_len, key_len)\n",
    "        # values shape: (N, value_len, heads, heads_dim)\n",
    "        # out after matrix multiply: (N, query_len, heads, head_dim), then\n",
    "        # we reshape and flatten the last two dimensions.\n",
    "\n",
    "        out = self.fc_out(out)\n",
    "        # Linear layer doesn't modify the shape, final shape will be\n",
    "        # (N, query_len, embed_size)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "05d5d74b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, embed_size, heads, dropout, forward_expansion):\n",
    "        super(TransformerBlock, self).__init__()\n",
    "        self.attention = SelfAttention(embed_size, heads)\n",
    "        self.norm1 = nn.LayerNorm(embed_size)\n",
    "        self.norm2 = nn.LayerNorm(embed_size)\n",
    "\n",
    "        self.feed_forward = nn.Sequential(\n",
    "            nn.Linear(embed_size, forward_expansion * embed_size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(forward_expansion * embed_size, embed_size),\n",
    "        )\n",
    "        # trace \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, value, key, query, mask):\n",
    "        attention = self.attention(value, key, query, mask)\n",
    "\n",
    "        # Add skip connection, run through normalization and finally dropout\n",
    "        x = self.dropout(self.norm1(attention + query))\n",
    "        forward = self.feed_forward(x)\n",
    "        out = self.dropout(self.norm2(forward + x))\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6bfb7406",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        src_vocab_size,\n",
    "        embed_size,\n",
    "        num_layers,\n",
    "        heads,\n",
    "        device,\n",
    "        forward_expansion,\n",
    "        dropout,\n",
    "        max_length,\n",
    "    ):\n",
    "\n",
    "        super(Encoder, self).__init__() \n",
    "        self.embed_size = embed_size\n",
    "        self.device = device\n",
    "        self.word_embedding = nn.Embedding(src_vocab_size, embed_size)\n",
    "        self.position_embedding = nn.Embedding(max_length, embed_size)\n",
    "        self.layers = nn.ModuleList(\n",
    "            [\n",
    "                TransformerBlock(\n",
    "                    embed_size,\n",
    "                    heads,\n",
    "                    dropout=dropout,\n",
    "                    forward_expansion=forward_expansion,\n",
    "                )\n",
    "                for _ in range(num_layers)\n",
    "            ]\n",
    "        )\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        N, seq_length = x.shape\n",
    "        positions = torch.arange(0, seq_length).expand(N, seq_length).to(self.device)\n",
    "        out = self.dropout(\n",
    "            (self.word_embedding(x) + self.position_embedding(positions))\n",
    "        )\n",
    "\n",
    "        # In the Encoder the query, key, value are all the same, it's in the\n",
    "        # decoder this will change. This might look a bit odd in this case.\n",
    "        for layer in self.layers:\n",
    "            out = layer(out, out, out, mask)\n",
    "\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d624184b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderBlock(nn.Module):\n",
    "    def __init__(self, embed_size, heads, forward_expansion, dropout, device):\n",
    "        super(DecoderBlock, self).__init__()\n",
    "        \n",
    "        self.norm = nn.LayerNorm(embed_size)\n",
    "        \n",
    "        self.attention = SelfAttention(embed_size, heads=heads)\n",
    "        self.transformer_block = TransformerBlock(\n",
    "            embed_size, heads, dropout, forward_expansion\n",
    "        )\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, value, key, src_mask, trg_mask):\n",
    "        attention = self.attention(x, x, x, trg_mask)\n",
    "        query = self.dropout(self.norm(attention + x))\n",
    "        out = self.transformer_block(value, key, query, src_mask)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6819b5e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        trg_vocab_size,\n",
    "        embed_size,\n",
    "        num_layers,\n",
    "        heads,\n",
    "        forward_expansion,\n",
    "        dropout,\n",
    "        device,\n",
    "        max_length,\n",
    "    ):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.device = device\n",
    "        self.word_embedding = nn.Embedding(trg_vocab_size, embed_size)\n",
    "        self.position_embedding = nn.Embedding(max_length, embed_size)\n",
    "\n",
    "        self.layers = nn.ModuleList(\n",
    "            [\n",
    "                DecoderBlock(embed_size, heads, forward_expansion, dropout, device)\n",
    "                for _ in range(num_layers)\n",
    "            ]\n",
    "        )\n",
    "        self.fc_out = nn.Linear(embed_size, trg_vocab_size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, enc_out, src_mask, trg_mask):\n",
    "        N, seq_length = x.shape\n",
    "        positions = torch.arange(0, seq_length).expand(N, seq_length).to(self.device)\n",
    "        x = self.dropout((self.word_embedding(x) + self.position_embedding(positions)))\n",
    "\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, enc_out, enc_out, src_mask, trg_mask)\n",
    "\n",
    "        out = self.fc_out(x)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "be91483b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        src_vocab_size :int,\n",
    "        trg_vocab_size :int,\n",
    "        src_pad_idx :int,\n",
    "        trg_pad_idx :int,\n",
    "        embed_size=512 :int,\n",
    "        num_layers=6 :int,\n",
    "        forward_expansion=4 :int,\n",
    "        heads=8 :int,\n",
    "        dropout=0 :int,\n",
    "        device=\"cpu\" :str,\n",
    "        max_length=100 :int,\n",
    "    ):\n",
    "\n",
    "        super(Transformer, self).__init__()\n",
    "\n",
    "        \n",
    "        self.encoder = Encoder(\n",
    "            src_vocab_size,\n",
    "            embed_size,\n",
    "            num_layers,\n",
    "            heads,\n",
    "            device,\n",
    "            forward_expansion,\n",
    "            dropout,\n",
    "            max_length,\n",
    "        )\n",
    "\n",
    "        self.decoder = Decoder(\n",
    "            trg_vocab_size,\n",
    "            embed_size,\n",
    "            num_layers,\n",
    "            heads,\n",
    "            forward_expansion,\n",
    "            dropout,\n",
    "            device,\n",
    "            max_length,\n",
    "        )\n",
    "\n",
    "        self.src_pad_idx = src_pad_idx\n",
    "        self.trg_pad_idx = trg_pad_idx\n",
    "        self.device = device\n",
    "\n",
    "    def make_src_mask(self, src):\n",
    "        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)\n",
    "        # (N, 1, 1, src_len)\n",
    "        return src_mask.to(self.device)\n",
    "\n",
    "    def make_trg_mask(self, trg):\n",
    "        N, trg_len = trg.shape\n",
    "        trg_mask = torch.tril(torch.ones((trg_len, trg_len))).expand(\n",
    "            N, 1, trg_len, trg_len\n",
    "        )\n",
    "\n",
    "        return trg_mask.to(self.device)\n",
    "\n",
    "    def forward(self, src, trg):\n",
    "        src_mask = self.make_src_mask(src)\n",
    "        trg_mask = self.make_trg_mask(trg)\n",
    "        enc_src = self.encoder(src, src_mask)\n",
    "        out = self.decoder(trg, enc_src, src_mask, trg_mask)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cb673eba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n",
      "torch.Size([2, 7, 10])\n"
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    print(device)\n",
    "\n",
    "    x = torch.tensor([[1, 5, 6, 4, 3, 9, 5, 2, 0], [1, 8, 7, 3, 4, 5, 6, 7, 2]]).to(\n",
    "        device\n",
    "    )\n",
    "    trg = torch.tensor([[1, 7, 4, 3, 5, 9, 2, 0], [1, 5, 6, 2, 4, 7, 6, 2]]).to(device)\n",
    "\n",
    "    src_pad_idx = 0\n",
    "    trg_pad_idx = 0\n",
    "    src_vocab_size = 10\n",
    "    trg_vocab_size = 10\n",
    "    model = Transformer(src_vocab_size, trg_vocab_size, src_pad_idx, trg_pad_idx, device=device).to(\n",
    "        device\n",
    "    )\n",
    "    out = model(x, trg[:, :-1])\n",
    "    print(out.shape)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
